#include "robot_sim_mi9/SimDriverMi9.hpp"

#include "rclcpp/node.hpp"
#include "webots/accelerometer.h"
#include "webots/camera.h"
#include "webots/gyro.h"
#include "webots/inertial_unit.h"
#include "webots/motor.h"
#include "webots/robot.h"

#include <functional>
#include <vector>

bool enableCamera = false;

inline int GetPeriod(int hz)
{
    return 1000 / hz;
}

void SimDriverMi9::init(webots_ros2_driver::WebotsNode *node, std::unordered_map<std::string, std::string> &params)
{
    (void)node;
    (void)params;
    leftMotor = wb_robot_get_device("WheelMotorLeft");
    rightMotor = wb_robot_get_device("WheelMotorRight");

    IMUAngle = wb_robot_get_device("IMUAngle");
    IMUGyro = wb_robot_get_device("IMUGyro");
    IMUAccelerometer = wb_robot_get_device("IMUAccelerometer");

    leftCamera = wb_robot_get_device("CameraLeft");
    rightCamera = wb_robot_get_device("CameraRight");

    wb_motor_set_torque(leftMotor, 0);
    wb_motor_set_torque(rightMotor, 0);
    // wb_motor_set_available_torque(leftMotor, 100);
    // wb_motor_set_available_torque(rightMotor, 100);

    int IMUPeriod = GetPeriod(200);
    wb_inertial_unit_enable(IMUAngle, IMUPeriod);
    wb_gyro_enable(IMUGyro, IMUPeriod);
    wb_accelerometer_enable(IMUAccelerometer, IMUPeriod);

    if (enableCamera)
    {
        wb_camera_enable(leftCamera, GetPeriod(30));
        wb_camera_enable(rightCamera, GetPeriod(30));
    }

    auto qos = rclcpp::SensorDataQoS().reliable();

    leftWheelPub = node->create_publisher<WheelSpeedMsg>("/left_motor_speed", qos);
    rightWheelPub = node->create_publisher<WheelSpeedMsg>("/right_motor_speed", qos);
    imuPub = node->create_publisher<IMUMsg>("/IMU", qos);

    leftCameraPub = node->create_publisher<CameraMsg>("/left_camera", qos);
    rightCameraPub = node->create_publisher<CameraMsg>("/right_camera", qos);

    leftMotorForceSub =
        node->create_subscription<MotorForceMsg>("/left_motor_force", rclcpp::SystemDefaultsQoS().reliable(),
                                                 [this](const MotorForceMsg &msg) { leftMotorForce = msg.data; });
    rightMotorForceSub =
        node->create_subscription<MotorForceMsg>("/right_motor_force", rclcpp::SystemDefaultsQoS().reliable(),
                                                 [this](const MotorForceMsg &msg) { rightMotorForce = msg.data; });
}

void SimDriverMi9::step()
{
    wb_motor_set_torque(leftMotor, leftMotorForce);
    wb_motor_set_torque(rightMotor, rightMotorForce);
    //  RCLCPP_INFO(rclcpp::get_logger("SimDriverMi9"),"Troque is %f:%f",leftMotorForce,rightMotorForce);

    WheelSpeedMsg leftSpeedMsg;
    leftSpeedMsg.data = wb_motor_get_velocity(leftMotor);
    leftWheelPub->publish(leftSpeedMsg);
    WheelSpeedMsg rightSpeedMsg;
    rightSpeedMsg.data = wb_motor_get_velocity(rightMotor);
    rightWheelPub->publish(rightSpeedMsg);
    RCLCPP_INFO(rclcpp::get_logger("SimDriverMi9"), "Speed is %f:%f", leftSpeedMsg.data, rightSpeedMsg.data);

    const double *quaternion = wb_inertial_unit_get_quaternion(IMUAngle);
    const double *gyro = wb_gyro_get_values(IMUGyro);
    const double *acc = wb_accelerometer_get_values(IMUAccelerometer);
    IMUMsg ImuMsg;
    ImuMsg.header.stamp = rclcpp::Clock().now();
    ImuMsg.header.frame_id = "IMU_frame";
    ImuMsg.orientation.x = quaternion[0];
    ImuMsg.orientation.y = quaternion[1];
    ImuMsg.orientation.z = quaternion[2];
    ImuMsg.orientation.w = quaternion[3];
    ImuMsg.angular_velocity.x = gyro[0];
    ImuMsg.angular_velocity.y = gyro[1];
    ImuMsg.angular_velocity.z = gyro[2];
    ImuMsg.linear_acceleration.x = acc[0];
    ImuMsg.linear_acceleration.y = acc[1];
    ImuMsg.linear_acceleration.z = acc[2];
    imuPub->publish(ImuMsg);

    if (!enableCamera)
        return;

    int width = wb_camera_get_width(leftCamera);
    int height = wb_camera_get_height(leftCamera);
    const unsigned char *img = wb_camera_get_image(leftCamera);
    int stepSize = width * 4;
    int totalSize = stepSize * height;
    std::vector<unsigned char> leftBytes(img, img + totalSize);
    CameraMsg leftCameraMsg;
    leftCameraMsg.header.stamp = rclcpp::Clock().now();
    leftCameraMsg.header.frame_id = "left_camera_frame";
    leftCameraMsg.height = height;
    leftCameraMsg.width = width;
    leftCameraMsg.encoding = "bgra8";
    leftCameraMsg.data = leftBytes;
    leftCameraMsg.is_bigendian = true;
    leftCameraMsg.step = stepSize;
    leftCameraPub->publish(leftCameraMsg);

    width = wb_camera_get_width(rightCamera);
    height = wb_camera_get_height(rightCamera);
    img = wb_camera_get_image(rightCamera);
    stepSize = width * 4;
    totalSize = stepSize * height;
    std::vector<unsigned char> rightBytes(img, img + totalSize);
    CameraMsg rightCameraMsg;
    rightCameraMsg.header.stamp = rclcpp::Clock().now();
    rightCameraMsg.header.frame_id = "right_camera_frame";
    rightCameraMsg.height = height;
    rightCameraMsg.width = width;
    rightCameraMsg.encoding = "bgra8";
    rightCameraMsg.data = rightBytes;
    rightCameraMsg.is_bigendian = true;
    rightCameraMsg.step = stepSize;
    rightCameraPub->publish(rightCameraMsg);
}

#include "pluginlib/class_list_macros.hpp"
PLUGINLIB_EXPORT_CLASS(SimDriverMi9, webots_ros2_driver::PluginInterface)
